package main

import (
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"sync"
)

var (
	upgrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
	connections = make(map[string]*websocket.Conn) // 用连接ID存储连接
	connMutex   sync.Mutex                         // 保护并发访问
)

func main() {
	r := gin.Default()
	r.GET("/ping", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"message": "pong",
		})
	})
	r.GET("/ws", wsHandler)
	r.GET("/test", ws)
	r.POST("/send", sendMessage)
	r.Run() // listen and serve on 0.0.0.0:8080 (for windows "localhost:8080")
}

func ws(c *gin.Context) {
	//Upgrade get request to webSocket protocol
	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer ws.Close()
	for {
		//read data from ws
		mt, message, err := ws.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}
		log.Printf("recv: %s", message)

		//write ws data
		err = ws.WriteMessage(mt, message)
		if err != nil {
			log.Println("write:", err)
			break
		}
	}
}

func wsHandler(c *gin.Context) {
	// 从查询参数获取连接ID，如果没有则生成一个
	connID := c.Query("id")
	if connID == "" {
		connID = "conn_" + c.ClientIP() // 简单用IP作为ID，实际应该用更唯一的方式
	}

	// Upgrade HTTP request to WebSocket
	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer ws.Close()

	// 注册新连接
	connMutex.Lock()
	connections[connID] = ws
	log.Printf("New connection registered. ID: %s, Total: %d", connID, len(connections))
	connMutex.Unlock()

	// 连接关闭时移除
	defer func() {
		connMutex.Lock()
		delete(connections, connID)
		log.Printf("Connection removed. ID: %s, Remaining: %d", connID, len(connections))
		connMutex.Unlock()
	}()

	for {
		// 读取客户端消息
		mt, message, err := ws.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}
		log.Printf("recv from %s: %s", connID, message)

		// 回显消息
		err = ws.WriteMessage(mt, message)
		if err != nil {
			log.Println("write:", err)
			break
		}
	}
}

// sendMessage 向特定连接发送消息
func sendMessage(c *gin.Context) {
	var request struct {
		ConnID  string `json:"conn_id"`
		Message string `json:"message"`
	}

	if err := c.ShouldBindJSON(&request); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request"})
		return
	}

	// 默认消息为"hello"
	message := request.Message
	if message == "" {
		message = "hello"
	}

	connMutex.Lock()
	ws, exists := connections[request.ConnID]
	connMutex.Unlock()

	if !exists {
		c.JSON(http.StatusNotFound, gin.H{"error": "Connection not found"})
		return
	}

	// 发送消息
	if err := ws.WriteMessage(websocket.TextMessage, []byte(message)); err != nil {
		log.Printf("Failed to send message to %s: %v", request.ConnID, err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to send message"})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"status":  "Message sent",
		"to":      request.ConnID,
		"message": message,
	})
}
